#!/usr/bin/env python
""" 
Basile Clout, September 2nd, 2007
pyWrekaMapd is a local deamon waiting for the client user to interrogate it
 Give:
     - direct neighbors: nmap
     - cpu relative power (pim2 10000)
     - latency and bandwidth with the nearest neighbors (using MPI: wreka_MPI)
 Interrogated by pyWrekaMapc.py, give the characteristics of the network corresponding to the local node

Command:
-h, --help Display this help
--cpu_on Enable cpu measurement
--cpu_off Disable cpu measurement
--set_user Give the user under which we measure cpu power
--net_on Enable network (bandwidth and latency) measurements
--net_off Disable network measurements
--set_list arg Set the list of available nodes in the cluster to arg
--set_nmap arg Give the absolute path to nmap
--set_nmapout arg Give the address of the nmap output
--set_wrekaMPI arg Give the absolute path to wreka_MPI
--set_pim2 arg Give the absolute path to pim2
--set_pim2size arg Determines the size of the pim2 test
--set_mpirun arg Give the absolute path to mpirun (to launch mpi)
--set_logfile arg Give the name of the logfile

Example:
lids04:~/wrekamap# ./pywrekamapd.py --net_off --cpu_off
node lids04 172.16.0.12 0 1 0 0
neighbors 172.16.0.12 172.16.0.14 172.16.0.17 172.16.0.10 
"""
 
import os
import string
import pprint
import time
import socket
from datetime import datetime
import pwd
import sys
import getopt


class node:
    """ Contains the networking characteristics of one node """
    
    def __init__(self, list, user, cpu_user):
        """ Initialize an instance of a node. 
        string list contains the addresses of all the node in the network
        """
        self.user = user
        self.cpu_user = cpu_user
        self.list = list
        self.neighbors = {}
        self.cpu = 0
        self.links_tasks ={}
        self.links_config = {}
        self.links_results = {}
        
        self.info = socket.gethostbyaddr(socket.gethostname())
        
    def getNeighbors(self, log, nmap_bin, nmap_out):
        """ Build the list of immediate neighbors of the node
        Print output in log
        tmp = directory containing temporary files
        """
        
        order = nmap_bin + " -send-ip -sP --ttl 0 " + self.list + " -oG " + nmap_out
        log.write("order sent: " + order + "\n")
        stdin, stdout, stderr = os.popen3(order)
        errput, output = stderr.read(), stdout.read()
        log.writelines("Output nmap: " + output +"\n")
        
        if errput != "":
            log.write("ERROR nmap:\n")
            log.writelines(errput)
            
        nmap = file(nmap_out, "r")
        
        for line in nmap:
            
            log.write(line)
            
            line.strip()
            if line[0] == "#":
                pass
            else:
                parts = line.split()
                address = parts[1]
                name = string.replace(parts[2], "(", "")
                name = string.replace(name, ")", "")
                
                if (name != self.info[0] or address != self.info[2][0]):
                    if name in self.list or address in self.list:
                        self.neighbors[name] = address 
                    
        log.write("neighbors dictionary:\n")
        pprint.pprint(self.neighbors, log)
        
        
    def getCPU(self, log, pim2_bin, pim2_size):
        """ Compute CPU time needed for computing pim2 10000"""
        
        order = pim2_bin + " " + pim2_size
        log.write(order + "\n")       

        save_euid = os.geteuid()
        try:
            uid = int(pwd.getpwnam(self.cpu_user)[2])
        except:
            log.write("ERROR: user " + self.cpu_user + " not found!\n")
#            print "Warning: User " + self.cpu_user + " not found!"
            self.cpu_user = self.user
        else:
            os.seteuid(uid)
            log.write("set uid to " + str(uid) + " for user " + self.cpu_user + "\n")

        self.cpu_user = pwd.getpwuid(os.geteuid())[0]
        #Here, we need the real (percepted) time, so we use time.time()
        t1 = time.time()
        stdin, stdout, stderr = os.popen3(order)
        stdout.read()
        self.cpu = time.time()-t1

        os.seteuid(save_euid)
        log.write("cpu time = " + str(self.cpu) + "\n")
    
    
    def buildHostfile(self, list, hosts_file, log):
        """ Build a hostfile suitable for use by the MPI implementation
        from the string list containing all the hosts (or ip addresses)
        into the file \"hosts_file\" 
        """
        log.write("Building hostfile ...")
        hostfile = file(hosts_file, "w")
        for line in list.split():
            hostfile.writelines(line + "\n")
        log.writelines("hostfile \"" + hosts_file + "\" written OK!\n")
        
                                                                           

    def getLinks(self, mpirun_bin, wrekaMPI, hosts_file, log):
        """ Basically, get the interesting values of the network by
        callin wrekaMPI and parsing its output
        """
        
        self.buildHostfile(self.list, hosts_file, log)
        try:
            master_node = self.list.split().index(self.info[0])
        except ValueError:
            try:
                master_node = self.list.split().index(self.info[2][0])
            except ValueError:
                log.write("ERROR: Local node " + self.info[0] + "/" + self.info[2][0] + " not in given list\n") 

        t = string.Template("sudo -H -u "+self.cpu_user+" $mpirun -machinefile $hostfile -n $number $wreka_MPI -m $master")
        order = t.substitute(mpirun = mpirun_bin, number = len(self.list.split()), hostfile = hosts_file, wreka_MPI = wrekaMPI, master = master_node)
        
        log.write("Launch wrekamap with :\n\t" + order + "\n")
        stdin, stdout, stderr = os.popen3(order)
        output, errput = stdout.read(), stderr.read()
        
        if errput != "":
            log.write("ERROR when launching wreka_MPI!:\n")
            log.writelines(errput)
            #sys.exit(1)
            
        else:
            log.write(" ... wreka_MPI launch OK!\n")
            
        log.write("Parse wreka_MPI output ...:\n")
        log.writelines(output + "\n")
        
        if (self.parse_wreka_output(output, log) == 1):
            log.write("Error!\n")
        else:
            log.write("OK!\n")
            log.write("links_config list:\n")
            pprint.pprint(self.links_config, log)
            log.write("links_tasks dictionary:\n")
            pprint.pprint(self.links_tasks, log)
            log.write("links_results dictionary:\n")
            pprint.pprint(self.links_results, log)
        
        
    def parse_wreka_output(self, output, log):    
        """ Parse the output of wrekaMPI, 
        and fills up the data structures of the node object 
        """
        
        next = 0 
        count = 0
        nbtasks = 0

        for line in output.splitlines():

            parts = line.split()

            if next == 0 and parts[0] == "config":
                next = 1                        # para "config"
            elif next == 0 and parts[0] == "tasks":
                next = 2                        # para "tasks"
            elif next == 0 and parts[0] == "results":
                next = 3                        # para "results"
            elif next == 0:
                log.write("Error during parsing. Pb with the paragraph's order\n")
            else:                                
                """ We are inside a paragraph """
                if next == 1:                    # in para "config"
                    if len(parts) != 11:
                        log.write("Error when parsing the \"config\" para\n")
                    else:
                        self.links_config["master"]=parts[0]
                        self.links_config["nbtasks"]=parts[1]
                        self.links_config["method"]=parts[2]
                        self.links_config["size"]=parts[3]
                        self.links_config["latsize"]=parts[4]
                        self.links_config["nbtests"]=parts[5]
                        self.links_config["skip"]=parts[6]
                        self.links_config["isrrequests"]=parts[7]
                        self.links_config["help"]=parts[8]
                        self.links_config["debug"]=parts[9]
                        self.links_config["live"]=parts[10]
                        
                        if self.links_config["help"]==1:
                            log.write("Warning: Help mode enabled!\n")
                        if self.links_config["debug"]==1:
                            log.write("Warning: Debug mode enabled!\n")
                        if self.links_config["live"]==1:
                            log.write("Info: Live mode enabled!\n")
                        
                        nbtasks = int(self.links_config["nbtasks"])
                        next = 0
                        
                elif next == 2:
                    if nbtasks == 0:
                        log.write("Error in paragraph order (tasks)\n")
                    if len(parts) != 3:
                        log.write("Error when parsing the \"tasks\" para\n")
                    self.links_tasks[parts[2]]=parts[:2]        # The key is the address (from self.neighbors)
                    count+=1
                    
                    if count == nbtasks:
                        next = 0
                        count = 0
                
                elif next == 3:
                    if nbtasks == 0:
                        log.write("Error in paragraph order (results)")
                    if len(parts) != 10:
                        log.write("Error when parsing the \"results\" para\n")
                    
                    self.links_results[parts[1]]=parts[2:]        # the key is the rank (from self.link_tasks)
                    count+=1
                    
                    if count == nbtasks:
                       next = 0
                       count = 0
              
    def printGrep(self, flags):
        """ the output of the deamon, in a grep-able way 
        flags contains the configuration: index 0 asks for cpu, 1 for network and 
        2 says there is an error
        TODO: XML output?
        """
        
        print "node", self.info[0], self.info[2][0], flags[0], flags[1], flags[2], flags[3]
        if flags[3] == 0:
            if flags[0] == 1:
                print "cpu", self.cpu, self.cpu_user
            if flags[1] == 1:
                neigh =  "neighbors " 
                for name, address in self.neighbors.iteritems():
                    neigh += address + " "
                print neigh
            if flags[2] == 1:
                for name, address in self.neighbors.iteritems():
                    out = "to "
                    rank = self.links_tasks[address][0]
                    mylist = self.links_results[rank]
                    out += address + " " + name
                    for el in mylist:
                        out+= " " + el
                    print out
                
        
def main():
    
    # Default configuration
    nmap_bin = "/usr/bin/nmap"
    pim2_bin = "/root/config/bin/bin/pim2"
    wrekaMPI_bin = "/home/basil/mpi/netMPI/wreka_MPI"
    mpirun_bin = "/usr/local/bin/mpiexec"

    pim2_size = "10000"
    cpu_user = "root"
    tmp_dir = "/tmp/wmapd/"
    hosts_file = tmp_dir + "hostfile.txt"
    log_file = "pywrekamapd.log"
    nmap_out = tmp_dir + "nmap_out.txt"

    list = "lids02 lids04 lids05 lids06 lids09 lids13 lids15 lids16"
    
    has_cpu = 1
    has_neigh = 1
    has_network = 1
    has_error = 0
    
    # Parse command line
    try:
        opts, args = getopt.getopt(sys.argv[1:],"h" ,["cpu_on", "cpu_off", "neigh_on", "neigh_off","net_on", "net_off", "set_list=", \
                                                   "set_nmap=", "set_wrekaMPI=", "set_pim2=", "set_nmapout=","set_pim2size=", "set_mpirun=", "set_logfile=","set_user=", "help"])
    except getopt.error, msg:
        print msg
        print "For help, use --help"
        sys.exit(2)
        
    # Process options
    for o, a in opts:
        if o in ("-h", "--help"):
            print __doc__
            sys.exit(0)
        elif o == "--cpu_on":
            has_cpu = 1
        elif o == "--cpu_off":
            has_cpu = 0
        elif o == "--neigh_on":
            has_neigh = 1
        elif o == "--neigh_off":
            has_neigh = 0
        elif o == "--net_off":
            has_network = 0
        elif o == "--net_on":
            has_network = 1
        elif o == "--set_list":
            list = a
        elif o == "--set_nmap":
            nmap_bin = a
        elif o == "--set_wrekaMPI":
            wrekaMPI_bin = a
        elif o == "--set_pim2":
            pim2_bin = a
        elif o == "--set_pim2size":
            pim2_size = a
        elif o == "--set_mpirun":
            mpirun_bin = a
        elif o == "--set_logfile":
            log_file = a
        elif o == "--set_user":
            cpu_user = a
        elif o == "--set_nmapout":
            nmap_out = a
        else:
            print "Not such option " + o
            sys.exit(1)

    
    # Initialize log file
    try:
        os.mkdir(tmp_dir)
    except OSError:
        print tmp_dir + " exists already\n"
            
    try:
        log = file(log_file, "w")
    except:
        print "Cannot create file " + log_file
        sys.exit(0)
        
    user = pwd.getpwuid(os.geteuid())[0] 
    log.write("Logfile for wrekamapd\nGenerated on " + str(datetime.now()) + " by " + user + "\n")
    
    log.write("\nLaunched with the command: " + string.join(sys.argv))
    log.write("\n\nnmap_bin = " + nmap_bin)
    log.write("\npim2_bin = " + pim2_bin)
    log.write("\nwrekaMPI_bin = " + wrekaMPI_bin)
    log.write("\nmpirun_bin = " + mpirun_bin)
    log.write("\npim2_size = " + pim2_size)
    log.write("\ntmp_dir = " + tmp_dir)
    log.write("\nhosts_file = " + hosts_file)
    log.write("\nlog_file = " + log_file)
    log.write("\nnmap_out = " + nmap_out)
    log.write("\nlist = " + list)
    log.write("\nuser = " + cpu_user)
    log.write("\nhas_cpu, has_neigh, has_network, has_error =  %d, %d, %d, %d\n\n" % (has_cpu, has_neigh, has_network, has_error))
    
    
    mynode = node(list, user, cpu_user)
    
    #CPU power informations?
    if (has_cpu == 1):
        try:
            os.stat(pim2_bin)
        except OSError:
            log.write("FATAL ERROR: Bad pim2 binary path?\n")
            has_error = 1
            #sys.exit(1)
        else:
            mynode.getCPU(log, pim2_bin, pim2_size)
       
    # Network informations?     
    if (has_neigh == 1):
        try:
            os.stat(nmap_bin)
        except OSError:
            log.write("FATAL ERROR: Bad nmap binary path?\n")
            #sys.exit(1)
        else:
            mynode.getNeighbors(log, nmap_bin, nmap_out)

    if (has_network == 1):
        try:
            os.stat(wrekaMPI_bin)
            os.stat(mpirun_bin)
        except OSError:
            log.write("FATAL ERROR: Bad wrekaMPI or mpirun binary path?\n")
            has_error = 1
            #sys.exit(1)
        else:
            mynode.getLinks(mpirun_bin, wrekaMPI_bin, hosts_file, log)
    
    flags = (has_cpu, has_neigh, has_network, has_error)
    log.write("flags: ")
    pprint.pprint(flags, log)
    mynode.printGrep(flags)
    
    

    try:
        os.remove(nmap_out)
        os.remove(hosts_file)
    except OSError:
        log.write("Some temp file impossible to delete")
    try:
        os.rmdir(tmp_dir)
    except OSError:
        print "Impossible to delete " + tmp_dir
    
    summary = ("succeded", "failed")[has_error]
    log.write("\n\nProgram " + summary +" at " + str(datetime.now()))
    log.close()
    
    
if __name__ == "__main__":
    main()
            
            
        
        
